import datetime
import hashlib
import random
import string
import re
from random import randint
from entity import ConnectionInfo

import pymysql

from entity import ColumnInfo

# 缓存
cache = {}
# 系统自带的schema
exclude_schema_list = ['information_schema', 'mysql', 'performance_schema', 'sys']


# 创建连接
def connect(ci: ConnectionInfo):
    try:
        s = ci.host + '|' + str(ci.port) + '|' + ci.username + '|' + ci.password
        h = hashlib.md5(s.encode()).hexdigest()
        if cache.get(h) is not None:
            return h
        conn = pymysql.connect(host=ci.host, port=ci.port, user=ci.username, password=ci.password,
                               database='information_schema')
        cache[h] = ci
        print('连接成功', h)
        conn.close()
        return h
    except Exception as e:
        print("连接时出现异常", e)
        return ''


# 获取连接
def get_connection(token: str, db_name: str):
    if cache.get(token) is None:
        print('token is deprecated')
        return None
    ci = cache[token]
    return pymysql.connect(host=ci.host, port=ci.port, user=ci.username, password=ci.password,
                           database=db_name)


# 获取schema列表
def get_schema_list(token: str):
    conn = get_connection(token, 'information_schema')
    if conn is None:
        return "连接失效"
    cur = conn.cursor()
    cur.execute('SELECT * FROM SCHEMATA')
    schema_list = []
    for row in cur.fetchall():
        schema_name = row[1]
        if schema_name not in exclude_schema_list:
            schema_list.append(schema_name)
    cur.close()
    conn.close()
    return schema_list


# 获取指定schema下的table列表
def get_table_list(token: str, schema: str):
    conn = get_connection(token, 'information_schema')
    if conn is None:
        return "连接失效"
    cur = conn.cursor()
    cur.execute("SELECT * FROM TABLES WHERE `TABLE_SCHEMA` = '{}'".format(schema))
    table_list = []
    for row in cur.fetchall():
        table_list.append((row[2], row[-1]))
    cur.close()
    conn.close()
    return table_list


# 获取指定schema的指定table的column列表
def get_column_list(token: str, schema: str, table: str):
    conn = get_connection(token, 'information_schema')
    if conn is None:
        return "连接失效"
    cur = conn.cursor()
    cur.execute("SELECT * FROM COLUMNS WHERE `TABLE_SCHEMA` = '{}' AND `TABLE_NAME` = '{}'".format(schema, table))
    column_list = []
    for row in cur.fetchall():
        data = {
            "table_catalog": row[0],
            "table_schema": row[1],
            "table_name": row[2],
            "column_name": row[3],
            "ordinal_position": row[4],
            "column_default": row[5],
            "is_nullable": row[6],
            "data_type": row[7],
            "character_maximum_length": row[8],
            "character_octet_length": row[9],
            "numeric_precision": row[10],
            "numeric_scale": row[11],
            "datetime_precision": row[12],
            "character_set_name": row[13],
            "collation_name": row[14],
            "column_type": row[15],
            "column_key": row[16],
            "extra": row[17],
            "privileges": row[18],
            "column_comment": row[19]
        }
        column_list.append(data)
    cur.close()
    conn.close()
    return column_list


# 数据库字符串类型
string_types = ['char', 'varchar', 'text', 'mediumtext', 'longtext']
# 数字类型
number_types = ['bit', 'bigint', 'decimal', 'double', 'float', 'int', 'smallint', 'tinyint']
# 日期类型
date_types = ['datetime', 'time', 'timestamp']


# 生成随机字符串
def gen_random_string(min_size: int, max_size: int, only: str, prefix: str, suffix: str):
    allowed_chars = string.ascii_letters + string.digits
    if only is not None:
        if only == 'letter':
            allowed_chars = string.ascii_letters
        if only == 'number':
            allowed_chars = string.digits
    ret = ''.join(random.choice(allowed_chars) for x in range(randint(min_size, max_size)))
    # 前缀，可以多个，逗号分隔
    if prefix is not None and prefix != '':
        if prefix.find(',') > 0:
            ps = prefix.split(',')
            ret = ps[randint(0, len(ps) - 1)] + ret
        else:
            ret = prefix + ret
    # 后缀，可以多个，逗号分隔
    if suffix is not None and suffix != '':
        if suffix.find(',') > 0:
            ss = suffix.split(',')
            ret = ss[randint(0, len(ss) - 1)] + ret
        else:
            ret = ret + suffix
    return ret


# 填充预设字符串
def fill_string(tmp_value: str):
    if tmp_value.find('[now]') > -1:
        return tmp_value.replace('[now]', datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S'))
    if tmp_value.find('[date]') > -1:
        return tmp_value.replace('[date]', datetime.datetime.now().strftime('%Y-%m-%d'))
    if tmp_value.find('[time]') > -1:
        return tmp_value.replace('[time]', datetime.datetime.now().strftime('%H:%M:%S'))
    return tmp_value


def gen_by_sql(token: str, schema: str, value_string: str, sql_col_map: dict, col_value_map: dict):
    tmp_string = str(value_string)
    for k in sql_col_map:
        v = sql_col_map[k]
        if len(v) > 6 and v[:6] == 'select':
            pattens = re.findall('\\${[a-zA-Z\d_]+}', v)
            if len(pattens) > 0:
                for pat in pattens:
                    cname = pat[2:-1]
                    cval = col_value_map.get(cname)
                    if cval is not None and cval[:6] != 'select':
                        v = v.replace(pat, cval)
            conn = get_connection(token, schema)
            if conn is None:
                break
            cur = conn.cursor()
            try:
                cur.execute(v)
                res = cur.fetchone()[0]
                if type(res) is datetime.datetime:
                    res = res.strftime('%Y-%m-%d %H:%M:%S')
                col_value_map[k] = str(res)
                sql_col_map[k] = 'null'
                tmp_string = tmp_string.replace('[' + k + ']', "'" + str(res) + "'")
            except Exception as e:
                print('error-003:', k, v, e)
                continue
            finally:
                cur.close()
                conn.close()
    chk_val = re.findall('\\${[a-zA-Z\d_]+}', tmp_string)
    if len(chk_val) > 0:
        gen_by_sql(token, schema, tmp_string, sql_col_map, col_value_map)
    return tmp_string


# 生成数据
def generate_data(token: str, count: int, column_list: list[ColumnInfo]):
    data_list = []
    ref_table_col_map = {}
    first_col = column_list[0]
    schema = first_col.table_schema
    table = first_col.table_name
    for i in range(count):
        values = []
        col_value_map = {}
        sql_col_map = {}
        for col in column_list:
            # 自增主键
            if col.column_key == 'PRI' and col.extra == 'auto_increment':
                continue
            # 可空字段
            if col.is_nullable == "YES" and col.set_null:
                col_value_map[col.column_name] = 'null'
                values.append('null')
                continue
            # 固定字段
            if col.fixed_string is not None and col.fixed_string != '':
                fs = col.fixed_string.split(',')
                ss = fs[randint(0, len(fs) - 1)].strip()
                if col.data_type in string_types or col.data_type in date_types:
                    ss = "'" + ss + "'"
                ss = fill_string(ss)
                col_value_map[col.column_name] = ss
                values.append(str(ss))
                continue
            # 关联指定表字段的值
            if col.ref_table_col is not None and col.ref_table_col != '':
                if ref_table_col_map.get(col.column_name) is None or len(ref_table_col_map[col.column_name]) == 0:
                    tmp_str = col.ref_table_col.split('.')
                    tmp_sql = 'select count(1) from ' + '.'.join(tmp_str[:-1])
                    conn = get_connection(token, schema)
                    if conn is None:
                        continue
                    cur = conn.cursor()
                    try:
                        cur.execute(tmp_sql)
                        tmp_data = ('.'.join(tmp_str[:-1]), tmp_str[-1], int(cur.fetchone()[0]))
                        ref_table_col_map[col.column_name] = tmp_data
                    except Exception as e:
                        print('error-001:', col.column_name, e)
                        continue
                    finally:
                        cur.close()
                        conn.close()
                conn = get_connection(token, schema)
                if conn is None:
                    continue
                cur = conn.cursor()
                try:
                    row_num = randint(0, ref_table_col_map[col.column_name][2] - 1)
                    tmp_col = ref_table_col_map[col.column_name]
                    tmp_sql = 'select `' + tmp_col[1] + '` from `' + tmp_col[0] + '` limit ' + str(row_num) + ',1'
                    cur.execute(tmp_sql)
                    ss = cur.fetchone()[0]
                    if type(ss) is datetime.datetime:
                        ss = ss.strftime('%Y-%m-%d %H:%M:%S')
                    ss = "'" + str(ss) + "'"
                    col_value_map[col.column_name] = ss
                    values.append(ss)
                except Exception as e:
                    print('error-002:', col.column_name, ref_table_col_map[col.column_name], e)
                    continue
                finally:
                    cur.close()
                    conn.close()
                continue
            # 如果字段类型是sql语句
            if col.sql is not None and col.sql != '' and col.sql.strip()[:6] == 'select':
                if sql_col_map.get(col.column_name) is None or len(sql_col_map[col.column_name]) == 0:
                    sql_col_map[col.column_name] = col.sql.strip()
                values.append('[' + col.column_name + ']')
                continue
            tmp_value = ''
            # 字符串类型
            if col.data_type in string_types:
                min_size = col.length if col.length is not None else int(col.character_maximum_length / 2)
                max_size = col.length if col.length is not None else col.character_maximum_length
                ss = gen_random_string(min_size=min_size,
                                       max_size=max_size,
                                       only=col.only,
                                       prefix=col.prefix_string,
                                       suffix=col.suffix_string)
                tmp_value = "'" + ss + "'"
                # 特定符转换
                tmp_value = fill_string(tmp_value)
            # 数字类型
            if col.data_type in number_types:
                if col.min is not None or col.max is not None:
                    col.min = col.min if col.min is not None else int(
                        min(int(col.max), pow(10, col.numeric_precision)) / 2)
                    col.max = col.max if col.max is not None else int(col.min * 2)
                    n = randint(int(col.min), min(int(col.max), pow(10, col.numeric_precision)))
                    tmp_value = str(n)
                else:
                    tmp_value = str(randint(1, pow(10, col.numeric_precision)))
                if col.scale is not None:
                    tmp_scale_val = randint(1, pow(10, int(col.scale)))
                    tmp_value = tmp_value + '.' + str(tmp_scale_val)
            # 日期类型
            if col.data_type in date_types:
                if col.min is not None or col.max is not None:
                    td = datetime.timedelta(days=30)
                    col.min = col.min if col.min is not None else datetime.datetime.now() - td
                    col.max = col.max if col.max is not None else col.min + td
                    min_time = datetime.datetime.strptime(col.min, '%Y-%m-%d %H:%M:%S').timestamp()
                    max_time = datetime.datetime.strptime(col.max, '%Y-%m-%d %H:%M:%S').timestamp()
                    d = datetime.date.fromtimestamp(randint(int(min_time), int(max_time)))
                    tmp_value = "'" + d.strftime('%Y-%m-%d %H:%M:%S') + "'"
                else:
                    d = datetime.date.fromtimestamp(
                        randint(int(datetime.datetime.now().timestamp()), 1667850638))
                    tmp_value = "'" + d.strftime('%Y-%m-%d %H:%M:%S') + "'"
                # 特定符转换
                tmp_value = fill_string(tmp_value)
            col_value_map[col.column_name] = tmp_value
            values.append(tmp_value)
        value_string = '(' + ','.join(values) + ')'
        # 所有字段生成后，检查是否有需要通过sql生成的字段
        if len(sql_col_map) > 0:
            value_string = gen_by_sql(token, schema, value_string, sql_col_map, col_value_map)
        data_list.append(value_string)
    if token is not None and token not in ['null', 'undefined', 'None', 'none']:
        cols = []
        for col in column_list:
            if col.column_key == 'PRI' and col.extra == 'auto_increment':
                continue
            cols.append('`' + col.column_name + '`')
        sql = 'insert into `' + table + '`(' + (','.join(cols)) + ') values ' + (','.join(data_list))
        conn = get_connection(token, schema)
        if conn is None:
            return "连接失效"
        cur = conn.cursor()
        try:
            cur.execute(sql)
        except Exception as e:
            print('error-004:', sql, e)
        conn.commit()
        cur.close()
        conn.close()
    ref_table_col_map.clear()
    return data_list
